package main

import (
	"context"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	media "github.com/livekit/media-sdk"
	lkpacer "github.com/livekit/mediatransportutil/pkg/pacer"
	lkauth "github.com/livekit/protocol/auth"
	livekitpb "github.com/livekit/protocol/livekit"
	lksdk "github.com/livekit/server-sdk-go/v2"
	lkmedia "github.com/livekit/server-sdk-go/v2/pkg/media"
	webrtc "github.com/pion/webrtc/v4"

	"github.com/hajimehoshi/oto/v2"
	"github.com/joho/godotenv"
)

// Simple PCM16 -> bytes helper
func int16ToBytes(samples []int16) []byte {
	out := make([]byte, len(samples)*2)
	for i := range samples {
		binary.LittleEndian.PutUint16(out[i*2:i*2+2], uint16(samples[i]))
	}
	return out
}

func main() {
	// Load .env from current folder if present
	_ = godotenv.Load()

	// Flags/env: LIVEKIT_URL, TOKEN, ROOM (optional), TARGET_IDENTITY filter
	var url string
	var token string
	var target string
	var room string
	var identity string
	var device string
	var inspect bool
	var watch bool
	var expectedSR int
	var beep bool
	var beepDur float64
	var pipeBeep bool
	flag.StringVar(&url, "url", os.Getenv("LIVEKIT_URL"), "LiveKit URL (wss://...) or env LIVEKIT_URL")
	flag.StringVar(&token, "token", os.Getenv("LIVEKIT_TOKEN"), "Access token or env LIVEKIT_TOKEN")
	flag.StringVar(&target, "target", os.Getenv("TARGET_IDENTITY"), "Subscribe only to this participant identity (optional)")
	flag.StringVar(&room, "room", os.Getenv("LIVEKIT_ROOM_NAME"), "Room name to join (defaults from LIVEKIT_ROOM_NAME)")
	flag.StringVar(&identity, "identity", os.Getenv("LIVEKIT_IDENTITY"), "Client identity (optional; autogenerated if empty)")
	flag.StringVar(&device, "device", "default", "Audio output device hint (ignored on macOS)")
	flag.BoolVar(&inspect, "inspect", false, "List rooms/participants via RoomService before connecting")
	flag.BoolVar(&watch, "watch", false, "Poll participants every 5s via RoomService and log")
	flag.IntVar(&expectedSR, "expected-sr", 48000, "Expected audio sample rate for throughput metrics (Hz)")
	flag.BoolVar(&beep, "beep", false, "Play a 1s 440Hz test tone on start to verify audio output")
	flag.Float64Var(&beepDur, "beep-dur", 1.0, "Test tone duration in seconds (with --beep)")
	flag.BoolVar(&pipeBeep, "pipe-beep", false, "Inject a 1s 440Hz test tone into the main player pipe")
	flag.Parse()

	if url == "" {
		log.Fatalf("url is required (flag --url or env LIVEKIT_URL)")
	}

	// If no token provided, mint one using API key/secret and room
	if token == "" {
		apiKey := os.Getenv("LIVEKIT_API_KEY")
		apiSecret := os.Getenv("LIVEKIT_API_SECRET")
		if apiKey == "" || apiSecret == "" {
			log.Fatalf("LIVEKIT_TOKEN not provided and LIVEKIT_API_KEY/SECRET missing; cannot mint token")
		}
		if room == "" {
			log.Fatalf("room is required to mint token (flag --room or env LIVEKIT_ROOM_NAME)")
		}
		if identity == "" {
			identity = fmt.Sprintf("speaker-%d", time.Now().UnixNano())
		}
		at := lkauth.NewAccessToken(apiKey, apiSecret)
		at.SetIdentity(identity)
		at.SetName(identity)
		at.AddGrant(&lkauth.VideoGrant{RoomJoin: true, Room: room})
		jwt, err := at.ToJWT()
		if err != nil {
			log.Fatalf("mint token: %v", err)
		}
		token = jwt
		log.Printf("joining as (minted): identity=%s room=%s", identity, room)
	} else {
		sub, name := decodeJWTIdentity(token)
		if sub != "" || name != "" {
			log.Printf("joining with provided token: sub=%s name=%s room=%s", sub, name, room)
		} else {
			log.Printf("joining with provided token (identity unknown pre-connect) room=%s", room)
		}
	}

	log.Printf("livekit-speaker: url=%s room=%s target=%q tokenLen=%d", url, room, target, len(token))

	// Optional: inspect via RoomService
	if inspect || watch {
		apiKey := os.Getenv("LIVEKIT_API_KEY")
		apiSecret := os.Getenv("LIVEKIT_API_SECRET")
		if apiKey == "" || apiSecret == "" {
			log.Printf("inspect/watch requested but LIVEKIT_API_KEY/SECRET missing; skipping REST")
		} else {
			apiURL := toHTTPBase(url)
			rs := lksdk.NewRoomServiceClient(apiURL, apiKey, apiSecret)
			// One-shot inspect
			if inspect {
				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				rooms, err := rs.ListRooms(ctx, &livekitpb.ListRoomsRequest{})
				if err != nil {
					log.Printf("ListRooms error: %v", err)
				} else {
					log.Printf("rooms: %d", len(rooms.Rooms))
					for _, r := range rooms.Rooms {
						log.Printf("- room: %s participants=%d", r.Name, r.NumParticipants)
					}
				}
				if room != "" {
					parts, err := rs.ListParticipants(ctx, &livekitpb.ListParticipantsRequest{Room: room})
					if err != nil {
						log.Printf("ListParticipants error: %v", err)
					} else {
						log.Printf("participants in %s: %d", room, len(parts.Participants))
						for _, p := range parts.Participants {
							log.Printf("- %s tracks: audio=%t video=%t", p.Identity, hasAudio(p), hasVideo(p))
						}
					}
				}
			}
			// Watch loop
			if watch && room != "" {
				go func() {
					ticker := time.NewTicker(5 * time.Second)
					defer ticker.Stop()
					for range ticker.C {
						ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
						parts, err := rs.ListParticipants(ctx, &livekitpb.ListParticipantsRequest{Room: room})
						cancel()
						if err != nil {
							log.Printf("watch: ListParticipants error: %v", err)
							continue
						}
						ids := make([]string, 0, len(parts.Participants))
						for _, p := range parts.Participants {
							ids = append(ids, p.Identity)
						}
						log.Printf("watch: participants=%v", ids)
					}
				}()
			}
		}
	}

	// Prepare audio output (48kHz mono, 16-bit)
	const outSampleRate = 48000
	const outChannels = 1
	ctx, ready, err := oto.NewContext(outSampleRate, outChannels, 2)
	if err != nil {
		log.Fatalf("oto.NewContext: %v", err)
	}
	<-ready

	// Optional: play a quick test tone to verify output path
	if beep {
		go func() {
			if err := playSine(ctx, 440.0, beepDur, outSampleRate, outChannels, 0.3); err != nil {
				log.Printf("beep error: %v", err)
			} else {
				log.Printf("beep played (%0.1fs at 440Hz)", beepDur)
			}
		}()
	}
	pr, pw := io.Pipe()
	player := ctx.NewPlayer(pr)
	player.Play()
	defer player.Close()
	if pipeBeep {
		go func() {
			// short delay to ensure player is running
			time.Sleep(200 * time.Millisecond)
			if err := writeSineTo(pw, 440.0, beepDur, outSampleRate, outChannels, 0.7); err != nil {
				log.Printf("pipe-beep error: %v", err)
			} else {
				log.Printf("pipe-beep played (%0.1fs at 440Hz)", beepDur)
			}
		}()
	}

	// Prepare subscribe handler BEFORE connecting so new tracks trigger immediately
	var cancelFns []context.CancelFunc
	pcw := &pcmWriter{w: pw, levelEvery: 50}
	pcw.metrics.expectedBps = expectedSR * outChannels * 2
	// periodic status
	go func() {
		t := time.NewTicker(2 * time.Second)
		defer t.Stop()
		for range t.C {
			pcw.mu.Lock()
			last := pcw.lastWrite
			bytes := pcw.totalBytes
			frames := pcw.frames
			pcw.mu.Unlock()
			ago := "never"
			if !last.IsZero() {
				ago = time.Since(last).Truncate(100 * time.Millisecond).String()
			}
			// also print throughput snapshot
			mbps, pct := pcw.metrics.snapshot()
			log.Printf("audio status: frames=%d totalBytes=%d lastWrite=%s ago, rate=%0.1f kB/s (%d%%)", frames, bytes, ago, mbps/1024.0, pct)
		}
	}()
	// 1s-rate printer
	go func() {
		t := time.NewTicker(1 * time.Second)
		defer t.Stop()
		for range t.C {
			kbps, pct := pcw.metrics.roll()
			log.Printf("audio throughput: %0.1f kB/s (%d%% of expected)", kbps/1024.0, pct)
		}
	}()
	trackCb := lksdk.ParticipantCallback{
		OnTrackSubscribed: func(track *webrtc.TrackRemote, publication *lksdk.RemoteTrackPublication, rp *lksdk.RemoteParticipant) {
			if track.Kind() != webrtc.RTPCodecTypeAudio {
				return
			}
			if target != "" && string(rp.Identity()) != target {
				log.Printf("skip non-target track from %s", rp.Identity())
				return
			}
			owner := string(rp.Identity())
			sid := string(publication.SID())
			log.Printf("subscribed: owner=%s sid=%s", owner, sid)
			// codec diagnostics
			c := track.Codec()
			log.Printf("track codec: mime=%s clock=%dHz channels=%d fmtp=%q", c.MimeType, c.ClockRate, c.Channels, c.SDPFmtpLine)

			// Use PCMRemoteTrack to get 48k mono samples for direct playout
			writer := &pcm16ToOtoWriter{pw: pcw}
			pcmTrack, err := lkmedia.NewPCMRemoteTrack(track, writer,
				lkmedia.WithTargetSampleRate(outSampleRate),
				lkmedia.WithTargetChannels(outChannels),
				lkmedia.WithHandleJitter(true),
			)
			if err != nil {
				log.Printf("PCMRemoteTrack err: %v", err)
				return
			}
			ctx, cancel := context.WithCancel(context.Background())
			cancelFns = append(cancelFns, cancel)
			go func() {
				<-ctx.Done()
				pcmTrack.Close()
			}()
		},
		OnTrackUnsubscribed: func(track *webrtc.TrackRemote, publication *lksdk.RemoteTrackPublication, rp *lksdk.RemoteParticipant) {
			log.Printf("track unsubscribed: %s from %s", publication.Name(), rp.Identity())
		},
		OnTrackMuted: func(pub lksdk.TrackPublication, p lksdk.Participant) {
			log.Printf("track muted: %s by %s", pub.Name(), p.Identity())
		},
		OnTrackUnmuted: func(pub lksdk.TrackPublication, p lksdk.Participant) {
			log.Printf("track unmuted: %s by %s", pub.Name(), p.Identity())
		},
	}

	// Connect to room with pacer for smoother playout
	pf := lkpacer.NewPacerFactory(
		lkpacer.LeakyBucketPacer,
		lkpacer.WithBitrate(512_000),
		lkpacer.WithMaxLatency(100*time.Millisecond),
	)
	roomCallback := &lksdk.RoomCallback{
		ParticipantCallback: trackCb,
		OnDisconnected:      func() { log.Printf("room disconnected") },
		OnReconnecting:      func() { log.Printf("room reconnecting...") },
		OnReconnected:       func() { log.Printf("room reconnected") },
		OnParticipantConnected: func(rp *lksdk.RemoteParticipant) {
			log.Printf("participant connected: %s", rp.Identity())
		},
		OnParticipantDisconnected: func(rp *lksdk.RemoteParticipant) {
			log.Printf("participant disconnected: %s", rp.Identity())
		},
		OnActiveSpeakersChanged: func(parts []lksdk.Participant) {
			ids := make([]string, 0, len(parts))
			for _, p := range parts {
				ids = append(ids, string(p.Identity()))
			}
			log.Printf("active speakers: %v", ids)
		},
	}
	lkRoom, err := lksdk.ConnectToRoomWithToken(url, token, roomCallback, lksdk.WithPacer(pf), lksdk.WithAutoSubscribe(true))
	if err != nil {
		log.Fatalf("connect: %v", err)
	}
	defer lkRoom.Disconnect()

	log.Printf("connected: local=%s remotes=%d", lkRoom.LocalParticipant.Identity(), len(lkRoom.GetRemoteParticipants()))

	// Enumerate existing participants and publications
	for _, rp := range lkRoom.GetRemoteParticipants() {
		id := string(rp.Identity())
		log.Printf("existing participant: %s (publications enumeration not available in this SDK version)", id)
	}

	// Handle shutdown
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	<-sig
	for _, c := range cancelFns {
		c()
	}
	log.Printf("shutting down")
	_ = pw.Close()
}

// pcmWriter writes PCM16 bytes into an oto.Player and logs level.
type pcmWriter struct {
	w          io.Writer
	frames     int
	levelEvery int
	totalBytes int64
	lastWrite  time.Time
	mu         sync.Mutex
	metrics    audioMetrics
}

func (p *pcmWriter) Write(b []byte) (int, error) {
	n, err := p.w.Write(b)
	if err == nil {
		p.mu.Lock()
		p.frames++
		f := p.frames
		p.totalBytes += int64(n)
		p.lastWrite = time.Now()
		p.mu.Unlock()
		p.metrics.addBytes(n)
		if p.levelEvery > 0 && f%p.levelEvery == 0 {
			// compute simple RMS on last chunk
			var sum int64
			for i := 0; i+1 < len(b); i += 2 {
				v := int16(binary.LittleEndian.Uint16(b[i : i+2]))
				if v < 0 {
					v = -v
				}
				sum += int64(v)
			}
			avg := float64(sum) / float64(len(b)/2)
			log.Printf("level avgAbs=%.0f", avg)
		}
	}
	return n, err
}

// pcm16ToOtoWriter adapts media.PCM16Writer -> oto.Player bytes.
type pcm16ToOtoWriter struct {
	pw *pcmWriter
}

func (w *pcm16ToOtoWriter) WriteSample(sample media.PCM16Sample) error {
	data := []int16(sample)
	if len(data) == 0 {
		return nil
	}
	// metrics: gaps and byte counts
	now := time.Now()
	w.pw.metrics.mu.Lock()
	if !w.pw.metrics.lastSampleTime.IsZero() {
		gap := now.Sub(w.pw.metrics.lastSampleTime)
		if gap > 250*time.Millisecond {
			log.Printf("audio gap detected: %s since last sample", gap.Truncate(10*time.Millisecond))
		}
	}
	w.pw.metrics.lastSampleTime = now
	w.pw.metrics.mu.Unlock()
	w.pw.metrics.addBytes(len(data) * 2)
	// Convert to bytes for oto
	return writeAll(w.pw, int16ToBytes(data))
}

func (w *pcm16ToOtoWriter) Close() error { return nil }

func writeAll(w interface{ Write([]byte) (int, error) }, b []byte) error {
	for len(b) > 0 {
		n, err := w.Write(b)
		if err != nil {
			if strings.Contains(err.Error(), "interrupted") {
				continue
			}
			return err
		}
		b = b[n:]
	}
	return nil
}

// writeSineTo writes PCM16 sine into a writer (same pipe as main player).
func writeSineTo(w io.Writer, freq float64, durSec float64, sampleRate int, channels int, volume float64) error {
	if durSec <= 0 {
		return nil
	}
	n := int(float64(sampleRate) * durSec)
	amp := int16(32767 * volume)
	// mono tone buffer
	buf := make([]int16, n)
	for i := 0; i < n; i++ {
		t := 2 * math.Pi * freq * float64(i) / float64(sampleRate)
		buf[i] = int16(float64(amp) * math.Sin(t))
	}
	// expand to channels
	var interleaved []int16
	if channels == 1 {
		interleaved = buf
	} else {
		interleaved = make([]int16, n*channels)
		for i := 0; i < n; i++ {
			for c := 0; c < channels; c++ {
				interleaved[i*channels+c] = buf[i]
			}
		}
	}
	// write in chunks to avoid blocking
	bytes := int16ToBytes(interleaved)
	const chunk = 4096
	pos := 0
	for pos < len(bytes) {
		end := pos + chunk
		if end > len(bytes) {
			end = len(bytes)
		}
		if _, err := w.Write(bytes[pos:end]); err != nil {
			return err
		}
		pos = end
	}
	return nil
}

// playSine generates a simple PCM16 sine wave and plays it via oto.
func playSine(ctx *oto.Context, freq float64, durSec float64, sampleRate int, channels int, volume float64) error {
	if durSec <= 0 {
		return nil
	}
	n := int(float64(sampleRate) * durSec)
	amp := int16(32767 * volume)
	// mono buffer
	buf := make([]int16, n)
	for i := 0; i < n; i++ {
		t := 2 * math.Pi * freq * float64(i) / float64(sampleRate)
		buf[i] = int16(float64(amp) * math.Sin(t))
	}
	// expand to requested channels (duplicate mono)
	var interleaved []int16
	if channels == 1 {
		interleaved = buf
	} else {
		interleaved = make([]int16, n*channels)
		for i := 0; i < n; i++ {
			for c := 0; c < channels; c++ {
				interleaved[i*channels+c] = buf[i]
			}
		}
	}
	// stream to player
	pr, pw := io.Pipe()
	player := ctx.NewPlayer(pr)
	player.Play()
	go func() {
		_ = writeAll(pw, int16ToBytes(interleaved))
		_ = pw.Close()
	}()
	// small wait for the buffer to flush
	time.Sleep(time.Duration(durSec*1000)*time.Millisecond + 100*time.Millisecond)
	_ = player.Close()
	return nil
}

// audioMetrics tracks per-second bytes vs expected rate.
type audioMetrics struct {
	mu             sync.Mutex
	expectedBps    int
	windowStart    time.Time
	windowBytes    int
	lastSampleTime time.Time
}

func (m *audioMetrics) addBytes(n int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.windowStart.IsZero() {
		m.windowStart = time.Now()
	}
	m.windowBytes += n
}

// roll prints last 1s window and resets.
func (m *audioMetrics) roll() (bytesPerSec float64, pct int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.windowStart.IsZero() {
		m.windowStart = time.Now()
		return 0, 0
	}
	now := time.Now()
	dur := now.Sub(m.windowStart)
	bytes := m.windowBytes
	// reset window
	m.windowStart = now
	m.windowBytes = 0
	if dur <= 0 {
		return 0, 0
	}
	bps := float64(bytes) / dur.Seconds()
	pct = 0
	if m.expectedBps > 0 {
		pct = int((bps / float64(m.expectedBps)) * 100.0)
		if pct > 999 {
			pct = 999
		}
		if pct < 0 {
			pct = 0
		}
	}
	return bps, pct
}

// snapshot shows current window rate without resetting.
func (m *audioMetrics) snapshot() (bytesPerSec float64, pct int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.windowStart.IsZero() {
		return 0, 0
	}
	dur := time.Since(m.windowStart)
	if dur <= 0 {
		return 0, 0
	}
	bps := float64(m.windowBytes) / dur.Seconds()
	pct = 0
	if m.expectedBps > 0 {
		pct = int((bps / float64(m.expectedBps)) * 100.0)
		if pct > 999 {
			pct = 999
		}
		if pct < 0 {
			pct = 0
		}
	}
	return bps, pct
}

// decodeJWTIdentity best-effort decodes a JWT and returns sub/name without verifying signature.
func decodeJWTIdentity(tok string) (sub, name string) {
	parts := strings.Split(tok, ".")
	if len(parts) < 2 {
		return "", ""
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return "", ""
	}
	var m map[string]any
	if err := json.Unmarshal(payload, &m); err != nil {
		return "", ""
	}
	if v, ok := m["sub"].(string); ok {
		sub = v
	}
	if v, ok := m["name"].(string); ok {
		name = v
	}
	return sub, name
}

func toHTTPBase(wsURL string) string {
	if strings.HasPrefix(wsURL, "wss://") {
		return "https://" + strings.TrimPrefix(wsURL, "wss://")
	}
	if strings.HasPrefix(wsURL, "ws://") {
		return "http://" + strings.TrimPrefix(wsURL, "ws://")
	}
	// assume already http(s)
	return wsURL
}

func hasAudio(p *livekitpb.ParticipantInfo) bool {
	for _, ti := range p.Tracks {
		if ti.Type == livekitpb.TrackType_AUDIO && ti.Muted == false {
			return true
		}
	}
	return false
}

func hasVideo(p *livekitpb.ParticipantInfo) bool {
	for _, ti := range p.Tracks {
		if ti.Type == livekitpb.TrackType_VIDEO && ti.Muted == false {
			return true
		}
	}
	return false
}
